
import cv2
import numpy as np
import matplotlib.pyplot as plt

def crop_image(img, x, y, width, height):
    # 读取图像

    # 检查图像是否成功读取
    if img is None:
        print(f"Error: Could not read image")
        return None

        # 使用NumPy切片来裁剪图像
    cropped_img = img[y:y + height, x:x + width]

    # 返回裁剪后的图像（如果需要的话）
    return cropped_img
def Visual_Sobel_gradient(image, sobelx, sobely):
    # 计算梯度的幅度
    magnitude, _ = cv2.cartToPolar(sobelx, sobely)

    # 归一化梯度幅度到0-255
    min_val, max_val = np.min(magnitude), np.max(magnitude)
    if max_val - min_val == 0:  # 避免除以零
        max_val = min_val + 1
    normalized_magnitude = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    # 使用归一化后的梯度幅度作为可视化图像
    gradient_visualization = normalized_magnitude

    return gradient_visualization


def visualize_gradient_magnitude(image, sobelx, sobely):
    """
    根据Sobel算子的输出对原始图像的梯度进行可视化处理。

    参数:
    image: numpy.ndarray, 原始灰度图像
    sobelx: numpy.ndarray, Sobel X方向的梯度
    sobely: numpy.ndarray, Sobel Y方向的梯度

    返回:
    matplotlib.figure.Figure, 包含颜色映射的梯度图像
    """
    # 计算梯度大小（幅度）
    gradient_magnitude = np.sqrt(sobelx ** 2 + sobely ** 2)

    # 归一化梯度大小到0-1之间
    gradient_magnitude_normalized = (gradient_magnitude - np.min(gradient_magnitude)) / (
                np.max(gradient_magnitude) - np.min(gradient_magnitude))

    # 创建一个新的图像用于可视化梯度大小
    gradient_image = gradient_magnitude_normalized * 255  # 映射到0-255范围
    gradient_image = gradient_image.astype(np.uint8)  # 转换为8位无符号整数

    # 绘制梯度图像
    fig, ax = plt.subplots()
    ax.imshow(gradient_image, cmap='viridis')  # 使用'viridis'或其他颜色映射
    ax.axis('off')  # 去掉坐标轴

    # 显示图像
    # plt.show()

    # 如果需要返回图像而不是显示，则注释掉plt.show()，并返回gradient_image
    return gradient_image

def draw_gradient_arrows(image, sobelx, sobely, arrow_density=0.05, arrow_length_scale=15, arrow_width=1):
    """
    根据Sobel算子的输出在原始灰度图像上绘制梯度箭头。

    参数:
    image: numpy.ndarray, 原始灰度图像
    sobelx: numpy.ndarray, Sobel X方向的梯度
    sobely: numpy.ndarray, Sobel Y方向的梯度
    arrow_density: float, 箭头绘制的密度（0到1之间）
    arrow_length_scale: float, 箭头长度的缩放因子
    arrow_width: float, 箭头线宽

    返回:
    matplotlib.figure.Figure, 包含箭头的图像
    """
    # 确保图像是灰度图
    if len(image.shape) != 2:
        raise ValueError("Input image must be grayscale.")

        # 计算梯度幅度和方向
    magnitude, angle = np.linalg.norm(np.dstack((sobelx, sobely)), axis=2), np.arctan2(sobely, sobelx)

    # 转换为角度制（如果angle是弧度制）
    angle = np.degrees(angle)

    # 创建一个和原始图像大小相同的空白图像用于可视化
    fig, ax = plt.subplots()
    ax.imshow(image, cmap='gray')

    # 遍历图像像素，绘制箭头
    height, width = image.shape
    for y in range(0, height, max(1, int(height * (1 - arrow_density)))):
        for x in range(0, width, max(1, int(width * (1 - arrow_density)))):
            # 根据梯度幅度设置阈值来决定是否绘制箭头
            if magnitude[y, x] > np.mean(magnitude) * 0.5:  # 可调整阈值
                # 计算箭头长度
                arrow_length = magnitude[y, x] * arrow_length_scale

                # 箭头方向（角度制）
                arrow_direction = angle[y, x]

                # 计算箭头的终点坐标
                dx = arrow_length * np.cos(np.radians(arrow_direction))
                dy = -arrow_length * np.sin(np.radians(arrow_direction))  # 注意y轴方向向下

                # 确保箭头在图像内
                if 0 <= x + dx < width and 0 <= y - dy < height:
                    # 绘制箭头
                    ax.annotate('', xy=(x, y), xytext=(x + dx, y - dy),
                                arrowprops=dict(facecolor='red', arrowstyle='-|>', lw=arrow_width),
                                textcoords="offset points")

                    # 显示图像
    ax.axis('off')  # 去掉坐标轴
    return fig
def main():
    #1. 读取图像 裁剪 灰度
    file_name = '../imgs/fishcamera.jpg'
    img = cv2.imread(file_name)  # 替换为你的图像路径
    croped_image = crop_image(img, 200,200,200,200)
    gray_img = cv2.cvtColor(croped_image, cv2.COLOR_BGR2GRAY)
    cv2.imshow('1 gray_img', gray_img)  # 归一化到0-255范围以便显示

    #2. sobel
    sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=5)

    #3. Canny + HoughLinesP
    edges = cv2.Canny(gray_img, 50, 150, apertureSize=3)
    cv2.imshow('2 edges', edges)  # 归一化到0-255范围以便显示
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 100, minLineLength=100, maxLineGap=10)
    print(lines)
    # lines_img = np.zeros_like(croped_image)
    # for line in lines:
    #     x1, y1, x2, y2 = line[0]
    #     cv2.line(lines_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    # cv2.imshow('HoughLinesP lines_img', lines_img)  # 归一化到0-255范围以便显示

    abs_sobelx = np.abs(sobelx)
    abs_sobely = np.abs(sobely)

    # 叠加（相加）水平和垂直边缘的绝对值
    sobel_combined = cv2.addWeighted(abs_sobelx, 0.5, abs_sobely, 0.5, 0)
    cv2.imshow('Sobel Combined', np.uint8(sobel_combined))  # 归一化到0-255范围以便显示
    visualize_img = visualize_gradient_magnitude(gray_img, sobelx, sobely)
    cv2.imshow('visualize',visualize_img)
    edges = cv2.Canny(visualize_img, 50, 150, apertureSize=3)
    cv2.imshow('Canny visualize',edges)

    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()